package cn.linkai.star.graph.core.service.impl;

import cn.hutool.core.util.StrUtil;
import cn.linkai.star.graph.comfyui.client.pojo.ComfyuiTask;
import cn.linkai.star.graph.core.service.RedisService;
import com.alibaba.fastjson2.JSON;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Service;

import java.time.Duration;


@Service
public class RedisServiceImpl implements RedisService {

    private final static String DISTRIBUTED_ID_KEY = "DISTRIBUTED_ID";
    private final static String TASK_KEY_PREFIX = "task_";
    private final static String DISTRIBUTED_QUEUE_KEY = "DISTRIBUTED_QUEUE";
    private final static String RUN_TASK_KEY = "run_task_";

    @Autowired
    private StringRedisTemplate redisTemplate;
    /**
     * 提交生图任务到redis队列中
     * @param comfyuiTask
     * @return
     */
    @Override
    public ComfyuiTask addQueueTask(ComfyuiTask comfyuiTask) {
        //1.获取分布式自增id
        Long idx = redisTemplate.opsForValue ().increment (DISTRIBUTED_ID_KEY);
        //2、设置回任务对象中
        comfyuiTask.setIndex (idx);
        //添加任务数据redis的String中
        redisTemplate.opsForValue ().set(TASK_KEY_PREFIX+comfyuiTask.getId (), JSON.toJSONString (comfyuiTask));

        //3.任务id压入队列zset
        redisTemplate.opsForZSet ().add (DISTRIBUTED_QUEUE_KEY,comfyuiTask.getId (),idx);
        return comfyuiTask;
    }

    /**
     * 从redis中获取序列号最小的一个的任务
     * @return
     */
    @Override
    public ComfyuiTask popQueueTask() {
        //1、判断这个对了是否有任务
        Long size = redisTemplate.opsForZSet ().size (DISTRIBUTED_QUEUE_KEY);
        if(size <= 0){
            return null;
        }

        //2、获取队列中的最新任务id
        ZSetOperations.TypedTuple<String> task = redisTemplate.opsForZSet ().popMin (DISTRIBUTED_QUEUE_KEY);
        if(task == null){
            return null;
        }
        String taskId = task.getValue ();

        //3、通过任务id获取任务数据
        String json = redisTemplate.opsForValue ().getAndDelete (TASK_KEY_PREFIX+taskId);
        if(StrUtil.isNotEmpty ( json)){
            return JSON.parseObject (json, ComfyuiTask.class);
        }
        return null;
    }

    /**
     * 添加已启动的任务到Redis缓存中
     *
     * @param promptId 任务提示ID，用作缓存的键前缀
     * @param task Comfyui任务对象，将被序列化后存储
     */
    @Override
    public void addStartedTask(String promptId, ComfyuiTask task) {
        redisTemplate.opsForValue ().set(RUN_TASK_KEY+promptId, JSON.toJSONString (task), Duration.ofMinutes (10));

    }

    /**
     * 根据提示ID获取已启动的任务信息
     *
     * @param promptId 提示ID，用于标识特定的任务
     * @return 返回对应的ComfyuiTask对象，如果未找到则返回null
     */
    @Override
    public ComfyuiTask getStartedTask(String promptId) {
        String json = redisTemplate.opsForValue ().get (RUN_TASK_KEY+promptId);
        if(StrUtil.isNotEmpty ( json)){
            return JSON.parseObject (json, ComfyuiTask.class);
        }
        return null;
    }

    /**
     * 判断redis队列中是否有任务
     * @return
     */
    @Override
    public boolean hasQueueTask() {
        return redisTemplate.opsForZSet ().size (DISTRIBUTED_QUEUE_KEY) > 0;
    }
}
